import json
import os
import os.path as osp

import cv2
import numpy as np
from tqdm import tqdm


def get_img_list(path):
    img_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith('.jpg') and f.find('_mask') == -1:
                img_list.append(osp.join(root, f))

    assert len(img_list) > 0, "没有图片！"
    return img_list


def get_mask(img, contours):
    mask = np.zeros(img.shape[:2], np.uint8)

    for points in contours:
        cv2.fillPoly(mask, [points], 255)

    return mask


if __name__ == '__main__':
    rootpth = r'data/shushu/train/No2'
    save_path = r'data/shushu/train_mask/No2'
    os.makedirs(save_path, exist_ok=True)

    img_list = get_img_list(rootpth)
    total_data = []
    for im_path in img_list:
        label_path = im_path.replace('.jpg', '.json')
        assert os.path.exists(label_path), "path is not exist"
        total_data.append([im_path, label_path])

    for i in tqdm(range(len(total_data))):
        img_path, anno_path = total_data[i]
        # dir_path = osp.dirname(img_path)
        base_name = osp.basename(img_path)
        mask_name = base_name.replace('.jpg', '_mask.png')
        mask_path = osp.join(save_path, mask_name)

        img = cv2.imread(img_path)
        # print(img.shape)

        with open(anno_path, 'r') as f:
            content = json.loads(f.read())


        contours = []
        for shape in content['shapes']:
            points = shape['points']
            contours.append(np.array(points, dtype=np.int64))

        label_mask = get_mask(img, contours)

        cv2.imwrite(mask_path, label_mask)
        # cv2.imshow('label_mask', label_mask)
        # cv2.waitKey()
